#!/usr/bin/python
# -*- coding: utf-8 -*-
import sys, os
help_msg = """
Script to train a GIZA++ alignment model.

Usage:
    {prog} ACTION GIZA_DIR MOSES_DIR CORPUS MODEL SRC_LANG TRG_LANG OUTPUT_DIR
Example of usage:
    {prog} prepare ~/cluster/giza/giza++ ~/cluster/moses ~/cluster/data/europarl_10k ibm2 fr en ~/cluster/europarl_10k_ibm2

ACTION: prepare (tokenization and truecasing), giza_prepare (vocabulary and coocurrences), mkcls, giza_run (giza alignment), all
GIZA_DIR: directory containing the 'giza', 'mkcls', 'plain2snt', and 'snt2cooc' binaries
MOSES_DIR: directory containing moses utility scripts (with a 'scripts' subdirectory)
CORPUS: path of the input corpus (without the extension)
MODEL: ibm1, ibm2, hmm, ibm3, ibm4, ibm5, ibm6
SRC_LANG: extension of the source side of the corpus (e.g., 'fr')
TRG_LANG: extension of the target side of the corpus (e.g., 'en')
OUTPUT_DIR: directory in which to write the giza alignment and tables (the preparation steps output files in the same directory as the corpus)
"""

prepare = """
rm -f {corpus}.tok* {corpus}.true* {corpus}.clean* {corpus}_truecase-model*
{moses}/scripts/tokenizer/tokenizer.perl -l {trg} < {corpus}.{trg} > {corpus}.tok.{trg} -threads 8
{moses}/scripts/tokenizer/tokenizer.perl -l {src} < {corpus}.{src} > {corpus}.tok.{src} -threads 8
{moses}/scripts/recaser/train-truecaser.perl --model {corpus}_truecase-model.{trg} --corpus {corpus}.tok.{trg}
{moses}/scripts/recaser/train-truecaser.perl --model {corpus}_truecase-model.{src} --corpus {corpus}.tok.{src}
{moses}/scripts/recaser/truecase.perl --model {corpus}_truecase-model.{trg} < {corpus}.tok.{trg} > {corpus}.true.{trg}
{moses}/scripts/recaser/truecase.perl --model {corpus}_truecase-model.{src} < {corpus}.tok.{src} > {corpus}.true.{src}
{moses}/scripts/training/clean-corpus-n.perl {corpus}.true {src} {trg} {corpus}.clean 1 80
"""

giza_prepare = """
rm -f {corpus}.{src}.vcb* {corpus}.{trg}.vcb* {corpus}.*.snt
{giza}/plain2snt {corpus}.clean.{src} {corpus}.clean.{trg} -vcb1 {corpus}.{src}.vcb -vcb2 {corpus}.{trg}.vcb -snt1 {corpus}.{src}_{trg}.snt -snt2 {corpus}.{trg}_{src}.snt
{giza}/snt2cooc {corpus}.{src}_{trg}.cooc {corpus}.{src}.vcb {corpus}.{trg}.vcb {corpus}.{src}_{trg}.snt
"""

mkcls = """
{giza}/mkcls -n2 -p{corpus}.clean.{trg} -V{corpus}.{trg}.vcb.classes
{giza}/mkcls -n2 -p{corpus}.clean.{src} -V{corpus}.{src}.vcb.classes
"""

giza_run = """
rm -rf {output}
mkdir -p {output}
{giza}/GIZA++ -s {corpus}.{src}.vcb -t {corpus}.{trg}.vcb -c {corpus}.{src}_{trg}.snt -coocurrencefile {corpus}.{src}_{trg}.cooc {model} -t1 5 -t2 5 -th 5 -t3 5 -t345 5 -t6 5 -o giza -outputpath {output}
"""

actions = dict(prepare=prepare, giza_prepare=giza_prepare, mkcls=mkcls, giza_run=giza_run)

def run_commands(commands, **parameters):
    commands = commands.strip().format(**parameters).split('\n')
    for cmd in commands:
        cmd = cmd.strip()
        if cmd:
            sys.stderr.write(cmd + '\n')
            os.system(cmd)

def model_args(model_str):
    model_num = ['ibm1', 'ibm2', 'hmm', 'ibm3', 'ibm4', 'ibm5', 'ibm6'].index(model_str)
    model_args = '-m1 {} -m2 {} -mh {} -m3 {} -m4 {} -m5 {} -m6 {}'.format(
            *([5] * (model_num + 1) + [0] * (6 - model_num)))
    return model_args

def main():
    try:
        action, giza, moses, corpus, model, src, trg, output = sys.argv[1:]
    except:
        sys.exit(help_msg.format(prog=sys.argv[0]))
    
    if action == 'all':
        commands = prepare + giza_prepare + (mkcls if model not in ['ibm1', 'ibm2'] else '') + giza_run
    else:
        commands = actions[action]
    
    model = model_args(model)

    parameters = dict(giza=giza, moses=moses, corpus=corpus, model=model, src=src, trg=trg, output=output)
    run_commands(commands, **parameters)

if __name__ == '__main__':
    main()

